package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"my_code/biligoo/other"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/gorilla/websocket"
	"github.com/valyala/fastjson"
)

type KEY struct {
	Name  string //备注名
	Admin bool   //是否为管理
	//MaxRoom        int    //是管理的话最大监控数目
	Password       string //密码
	ExpireTime     int64  //超过这个时间戳拒绝连接
	MaxConnections int    //超过最大连接数拒绝连接
	Rand           int    //0-100的概率推送
}

func sendMessage(code int, message string) []byte {
	jsonResponse, err := json.Marshal(map[string]interface{}{"code": 0, "type": "message", "message": message})
	if err != nil {
		return []byte{}
	}
	return jsonResponse
}

var (
	errAccess  = sendMessage(1, "Key错误")
	errTimeOut = sendMessage(2, "Key过期")
	errConnNum = sendMessage(3, "Key连接已满")
	errFast    = sendMessage(4, "请求过快")
)

func main() {
	go func() {
		err := http.ListenAndServe(":56666", nil) //设置监听的端口
		if err != nil {
			fmt.Printf("%v 监听端口失败: %v\n", TI(), err)
		}
	}()
	go func() {
		for {

			rece := <-chMsg

			// 使用fastjson解析JSON
			p := fastjson.Parser{}
			value, err := p.Parse(rece)
			if err != nil {
				fmt.Println("Error parsing JSON:", err)
				continue
			}
			giftID := value.GetInt("raffle_id")
			if !giftIDs.Add(giftID) {
				raffleType := string(value.GetStringBytes("raffle_type"))

				// 获取start_time字段的值
				startTime := value.GetInt("start_time")
				if startTime == 0 {
					value.Set("start_time", fastjson.MustParse(fmt.Sprintf("%d", nowTime())))
					// 将修改后的value转换回JSON字符串
					modifiedJSON := value.MarshalTo(nil)
					rece = string(modifiedJSON)
				}

				var data map[string]interface{}
				err = json.Unmarshal([]byte(rece), &data)
				if err != nil {
					fmt.Println("json data err:", err)
					continue
				}
				b, _ := json.Marshal(map[string]interface{}{"code": 0, "type": "raffle", "data": data})
				fmt.Printf("%v推送抽奖ID: %v 类型: %v\n", TI(), giftID, raffleType)
				oldMsgRefresh(string(b))
				broadcast(b, false)
			}
			if giftID == 0 {

			}
		}
	}()
	for {
		l := ""
		scanner := bufio.NewScanner(os.Stdin)
		if !scanner.Scan() {
			continue
		}
		l = scanner.Text()
		switch l {
		case "":
			keys, keysAdmin := map[string][]string{}, map[string][]string{}
			mutex.Lock()
			for conn, key := range clients {
				var ips []string
				var ok bool
				if key.Admin {
					if ips, ok = keysAdmin[key.Password]; !ok {
						ips = []string{}
					}
					ips = append(ips, conn.RemoteAddr().String())
					keysAdmin[key.Password] = ips
				} else {
					if ips, ok = keys[key.Password]; !ok {
						ips = []string{}
					}
					ips = append(ips, conn.RemoteAddr().String())
					keys[key.Password] = ips
				}
			}
			mutex.Unlock()
			txt := "\r======连接情况======"
			for k, v := range keys {
				txt += fmt.Sprintf("\n\rKEY %v:", k)
				for kk, vv := range v {
					txt += fmt.Sprintf("\n\r  %v. %v", kk+1, vv)
				}
			}
			if len(keysAdmin) != 0 {
				txt += "\n"
			}
			for k, v := range keysAdmin {
				txt += fmt.Sprintf("\n\rKEY %v(管理):", k)
				for kk, vv := range v {
					monitorInfosLock.RLock()
					monitor, ok := monitorInfos[vv]
					monitorInfosLock.RUnlock()
					if ok {
						txt += fmt.Sprintf("\n  %v. %v (%v/%v-%v)", kk+1, vv, monitor.runs, monitor.not_runs+monitor.runs, monitor.MaxRoom)
					} else {
						txt += fmt.Sprintf("\n  %v. %v", kk+1, vv)
					}
				}
			}
			fmt.Println(txt)
		default:
			ls := strings.Split(l, " ")
			if len(ls) == 3 && ls[0] == "b" {
				fmt.Println(ls)
				if t, err := strconv.Atoi(ls[2]); err == nil {
					BlackIP(ls[1], nowTime()+60*60*int64(t))
					continue
				}
			}
			if len(ls) == 2 && ls[0] == "r" {
				fmt.Println(ls)
				ip := ls[1]
				connectionsLock.Lock()
				connectionsCounter[ip] = [maxConnectionsPerIP]int64{}
				connectionsLock.Unlock()
				continue
			}
			broadcast(sendMessage(0, l), true)
		}
	}
}

var chMsg chan string

type ConcurrentSet struct {
	sync.RWMutex
	set map[int]bool
}

func NewConcurrentSet() *ConcurrentSet {
	return &ConcurrentSet{
		set: make(map[int]bool),
	}
}

func (cs *ConcurrentSet) Add(num int) bool {
	cs.Lock()
	defer cs.Unlock()

	if _, ok := cs.set[num]; ok {
		return true
	}

	cs.set[num] = true
	return false
}

var giftIDs *ConcurrentSet

func init() {
	giftIDs = NewConcurrentSet()
	chMsg = make(chan string)

}

var (
	clients = make(map[*websocket.Conn]KEY) // connected clients
	mutex   = &sync.Mutex{}                 // mutex for map access
)

func init() {
	http.Handle("/", http.FileServer(http.Dir("./files")))
	http.HandleFunc("/ws", handleWebSocket)
	// http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
	// 	if r.URL.Path != "/" {
	// 		http.NotFound(w, r)
	// 		return
	// 	}
	// 	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	// 	fmt.Fprint(w, returnIndex())
	// })
}
func haveKey(a string, s []KEY) (have bool, key KEY) {
	key = KEY{}
	for _, v := range s {
		if a == v.Password {
			have = true
			key = v
		}
	}
	return
}

const maxConnectionsPerIP = 10
const banThreshold = 60 * 60 //maxConnectionsPerIP*5 + 1

var (
	connectionsCounter = make(map[string][maxConnectionsPerIP]int64)
	connectionsLock    sync.RWMutex
)
var (
	blackIP     = map[string]int64{}
	blackIPLock sync.RWMutex
)

func BlackIP(ip string, endTime int64) int64 {
	blackIPLock.Lock()
	defer blackIPLock.Unlock()
	if endTime == 0 {
		if t, ok := blackIP[ip]; ok {
			return t
		}
		return 0
	}
	blackIP[ip] = endTime
	t := time.Unix(endTime, 0).In(time.FixedZone("CST", 28800)).Format("01-02 15:04:05")
	fmt.Printf("%v拉黑IP: %v 至 %v\n", TI(), ip, t)
	return endTime
}
func nowTime() int64 {
	return time.Now().Unix()
}

var (
	oldMsg     = []string{}
	oldMsgLock sync.RWMutex
)

func oldMsgRefresh(msg string) {
	oldMsgLock.Lock()
	defer oldMsgLock.Unlock()
	oldMsg = append(oldMsg, msg)
	newMsg := []string{}
	for _, v := range oldMsg {
		p := fastjson.Parser{}
		value, err := p.Parse(v)
		if err != nil {
			fmt.Println("Error parsing JSON:", err)
			continue
		}
		//fmt.Println(int64(value.GetInt("data", "start_time")+value.GetInt("data", "time")), "?", nowTime())
		if int64(value.GetInt("data", "start_time")+value.GetInt("data", "time")) >= nowTime() {
			newMsg = append(newMsg, v)
		}
	}
	oldMsg = newMsg
}
func getOldMsg() []string {
	oldMsgLock.RLock()
	defer oldMsgLock.RUnlock()
	return oldMsg
}
func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	ip, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}
	now := time.Now().Unix()

	var timestamps [maxConnectionsPerIP]int64
	var ok bool

	connectionsLock.RLock()
	timestamps, ok = connectionsCounter[ip]
	connectionsLock.RUnlock()

	if !ok {
		connectionsLock.Lock()
		connectionsCounter[ip] = [maxConnectionsPerIP]int64{}
		timestamps = connectionsCounter[ip]
		connectionsLock.Unlock()
	}

	// Upgrade initial GET request to WebSocket
	upgrader := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		//log.Println(err)
		return
	}
	if now-timestamps[0] < banThreshold {
		writeMessage(conn, errFast)
		time.Sleep(time.Second * 5)
		//closeConn(conn)
		//return
	}

	connectionsLock.Lock()
	copy(timestamps[:maxConnectionsPerIP-1], timestamps[1:])
	timestamps[maxConnectionsPerIP-1] = now
	connectionsCounter[ip] = timestamps
	connectionsLock.Unlock()
	// 处理WebSocket连接逻辑
	// ...
	hello := func(key KEY, num int) (close bool) {
		msg := key.Name
		if key.ExpireTime != -1 {
			if msg != "" {
				msg += ", "
			}
			msg += fmt.Sprintf("%v到期", time.Unix(key.ExpireTime, 0).In(time.FixedZone("CST", 28800)).Format("2006-01-02 15:04:05"))
		}
		if key.MaxConnections != -1 {
			if msg != "" {
				msg += ", "
			}
			msg += fmt.Sprintf("%v/%v连接", num+1, key.MaxConnections)
		}
		if key.Rand != 100 {
			if msg != "" {
				msg += ", "
			}
			msg += fmt.Sprintf("%v%%推送", key.Rand)
		}
		err := writeMessage(conn, sendMessage(0, msg))
		if err != nil {
			close = true
			return
		}
		if key.Rand == 100 && nowTime()-timestamps[maxConnectionsPerIP-2] >= 10 {
			oldmsg := getOldMsg()
			for _, v := range oldmsg {
				err = writeMessage(conn, []byte(v))
				if err != nil {
					close = true
					return
				}
			}
		}
		return
	}
	// Get client password from query parameter ?password=HiBjYwmRMrvUIu9fsVkkmXPL2CeUNUwK
	password := r.URL.Query().Get("key")
	maxroom := r.URL.Query().Get("maxroom")
	if have, key := haveKey(password, TOML.Keys); have {
		if key.ExpireTime != -1 && key.ExpireTime < time.Now().Unix() {
			writeMessage(conn, errTimeOut)
			time.Sleep(time.Second * 5)
			closeConn(conn)
			return
		}
		num := 0
		isAdmin := key.Admin
		mutex.Lock()
		for _, v := range clients {
			if v.Password == key.Password {
				num++
			}
		}
		mutex.Unlock()
		if key.MaxConnections != -1 && key.MaxConnections <= num {
			writeMessage(conn, errConnNum)
			time.Sleep(time.Second * 5)
			closeConn(conn)
			return
		}
		// Lock the clients map before adding new connection
		close := hello(key, num)
		if close {
			closeConn(conn)
			return
		}
		mutex.Lock()
		clients[conn] = key
		mutex.Unlock()
		// Start a new goroutine for the connection
		go handleConnection(conn, isAdmin, maxroom)

		// Start a new goroutine for the heartbeat mechanism
		//go heartbeat(conn)
	} else {
		writeMessage(conn, errAccess)
		fmt.Printf("%v%v 尝试使用错误Key(%v)连接\n", TI(), conn.RemoteAddr().String(), password)
		time.Sleep(time.Second * 5)
		closeConn(conn)
	}

}

func handleConnection(conn *websocket.Conn, isAdmin bool, maxroom string) {
	defer func() {
		mutex.Lock()
		delete(clients, conn)
		mutex.Unlock()
		closeConn(conn)
	}()
	name := conn.RemoteAddr().String()
	if isAdmin {
		// 设置读取超时时间为 40 秒
		conn.SetReadDeadline(time.Now().Add(40 * time.Second))
		name += " [管理]"
		max, _ := strconv.Atoi(maxroom)
		monitorInfosLock.Lock()
		monitorInfos[conn.RemoteAddr().String()] = monitorInfo{conn, max, []int{}, []int{}, []int{}, 0, 0}
		monitorInfosLock.Unlock()
	}
	fmt.Printf("%v%v 连接到服务器\n", TI(), name)
	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			//log.Println(err)
			return
		}
		if isAdmin {
			p := fastjson.Parser{}
			js, err := p.Parse(string(message))
			if err != nil {
				fmt.Println("Error parsing JSON:", err)
				continue
			}
			if js.Exists("err") {
				fmt.Printf("%v%v发来: %v\n", TI(), name, string(message))
				continue
				num := js.GetInt("err")
				now := js.GetInt("now")
				monitorInfosLock.Lock()
				monitor := monitorInfos[conn.RemoteAddr().String()]
				monitor.MaxRoom = now - num
				if monitor.MaxRoom < 0 {
					monitor.MaxRoom = 0
				}
				fmt.Printf("%v监控 %v 最大房间数改为%v   \n", other.TI(), conn.RemoteAddr().String(), monitor.MaxRoom)
				monitorInfos[conn.RemoteAddr().String()] = monitor
				monitorInfosLock.Unlock()
			} else if js.Exists("runs") && js.Exists("not_runs") {
				runs := js.GetArray("runs")
				not_runs := js.GetArray("not_runs")
				//fmt.Printf("%v监控 %v 运行%v 未运行%v   \n", other.TI(), conn.RemoteAddr().String(), len(runs), len(not_runs))

				monitorInfosLock.Lock()
				monitor := monitorInfos[conn.RemoteAddr().String()]
				monitor.runs = len(runs)
				monitor.not_runs = len(not_runs)
				monitorInfos[conn.RemoteAddr().String()] = monitor
				monitorInfosLock.Unlock()
			} else if js.Exists("refresh") {
				var data map[string]interface{}
				err = json.Unmarshal([]byte(js.Get("refresh").String()), &data)
				if err != nil {
					fmt.Println("json data err:", err)
					continue
				}
				b, _ := json.Marshal(map[string]interface{}{"code": 0, "type": "refresh", "data": data})
				broadcast(b, false)
			} else if js.Exists("RoomGift") {
				var data map[string]interface{}
				err = json.Unmarshal([]byte(js.Get("RoomGift").String()), &data)
				if err != nil {
					fmt.Println("json data err:", err)
					continue
				}
				b, _ := json.Marshal(map[string]interface{}{"code": 0, "type": "RoomGift", "data": data})
				broadcast(b, false)
			} else if !js.Exists("raffle_id") {
				fmt.Printf("%v%v发来: %v\n", TI(), name, string(message))
			} else {
				chMsg <- string(message)
			}
			// 设置读取超时时间为 40 秒
			conn.SetReadDeadline(time.Now().Add(40 * time.Second))
		} else {
			fmt.Printf("%v%v发来: %v\n", TI(), name, string(message))
		}
	}
}
func getClientIP(conn *websocket.Conn) string {
	// 获取远程地址
	remoteAddr := conn.RemoteAddr()
	tcpAddr, ok := remoteAddr.(*net.TCPAddr)

	if !ok {
		return ""
	}

	// 获取IP地址
	ip := tcpAddr.IP.String()
	return ip
}

// 创建一个映射，用于存储每个conn对应的互斥锁
var connMutexMap = &sync.Map{}

// writeMessage函数
// func writeMessage(conn *websocket.Conn, message []byte) error {
// 	fmt.Printf("%v将写入 %v\n", TI(), conn.RemoteAddr().String())
// 	defer fmt.Printf("%v已写入 %v\n", TI(), conn.RemoteAddr().String())
// 	// 获取当前conn对应的互斥锁，如果不存在则创建新的互斥锁
// 	connMutex, _ := connMutexMap.LoadOrStore(conn, &sync.Mutex{})
// 	mutex := connMutex.(*sync.Mutex)

// 	// 使用conn对应的互斥锁保护写入操作
// 	mutex.Lock()
// 	defer mutex.Unlock()

// 	return conn.WriteMessage(websocket.TextMessage, message)
// }

// writeMessage function
func writeMessage(conn *websocket.Conn, message []byte) error {

	// Get the mutex for the current conn. If it doesn't exist, create a new one.
	connMutex, _ := connMutexMap.LoadOrStore(conn, &sync.Mutex{})
	mutex := connMutex.(*sync.Mutex)

	// Use the mutex for the current conn to protect the write operation.
	mutex.Lock()
	defer mutex.Unlock()

	// Set a timeout for the write operation.
	timeout := time.Duration(time.Second * 2)
	err := conn.SetWriteDeadline(time.Now().Add(timeout))
	if err != nil {
		return err
	}
	return conn.WriteMessage(websocket.TextMessage, message)
}

func closeConn(conn *websocket.Conn) {
	ip := conn.RemoteAddr().String()
	conn.Close()
	fmt.Printf("%v%v 断开连接\n", TI(), ip)
	monitorInfosLock.Lock()
	if _, ok := monitorInfos[ip]; ok {
		delete(monitorInfos, ip)
	}
	monitorInfosLock.Unlock()
}
func broadcast(message []byte, isSystem bool) {
	//fmt.Printf("%v推送: %v\n", TI(), string(message))
	rand.Seed(time.Now().UnixNano())
	randomNumber := rand.Intn(100) + 1
	// Iterate over all connected clients and write the message to their WebSocket
	for client, key := range clients {
		haveKey := false
		for _, k := range TOML.Keys {
			if k.Password == key.Password {
				haveKey = true
				if time.Now().Unix() < k.ExpireTime || k.ExpireTime == -1 {
					ip := getClientIP(client)
					//fmt.Println("ip", ip)
					if (randomNumber <= k.Rand || isSystem) && BlackIP(ip, 0) < nowTime() {
						err := writeMessage(client, message)
						if err != nil {
							//log.Println(err)
							closeConn(client)
							mutex.Lock()
							delete(clients, client)
							mutex.Unlock()
							continue
						}
					} else {
						if !k.Admin {
							//fmt.Printf("%vKEY: %v IP: %v 不推送", TI(), k.Password, ip)
						}
					}

				} else {
					writeMessage(client, errTimeOut)
					closeConn(client)
					mutex.Lock()
					delete(clients, client)
					mutex.Unlock()
				}
			}
		}
		if !haveKey {
			writeMessage(client, errAccess)
			closeConn(client)
			mutex.Lock()
			delete(clients, client)
			mutex.Unlock()
		}
	}
}

func heartbeat(conn *websocket.Conn) {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			// Send a ping message to the client to check if the connection is still alive
			err := conn.WriteMessage(websocket.PingMessage, nil) //////panic: concurrent write to websocket connection为什么这里会报错，怎么解决
			if err != nil {
				//log.Println(err)
				return
			}
		}
	}
}

func TI() string {
	return time.Now().In(time.FixedZone("CST", 28800)).Format("\r[01-02 15:04:05] ")
}

var TOML *To

func init() {
	go func() {
		for {
			TOML = ReadToml()
			time.Sleep(time.Second * 10)
		}
	}()
}

var (
	tomlFilePath = "./server-keys.toml"
)

type To struct {
	Keys []KEY
}

// WriteToml ...
func (T *To) WriteToml() {
	f, err := os.Create(tomlFilePath)
	if err != nil {
		fmt.Println(err)
		return
	}
	if err := toml.NewEncoder(f).Encode(T); err != nil {
		fmt.Println(err)
		return
	}
}
func ReadToml() (config *To) {
	if _, err := os.Stat(tomlFilePath); err != nil {
		s := To{Keys: []KEY{{"管理", true, "375ecc98e7d7629f5b816f66160184a0", -1, 100, 0}, {"测试", false, "", -1, 100, 80}}}
		s.WriteToml()
	}
	if _, err := toml.DecodeFile(tomlFilePath, &config); err != nil {
		fmt.Println(other.TI() + "请检查配置文件！")
		time.Sleep(time.Second)
		panic(err)
	}
	return
}

// 模拟从文件中读取内容的函数
func returnIndex() string {
	content, err := ioutil.ReadFile("index.html")
	if err != nil {
		return ""
	}
	return string(content)
}
